#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import re
import io
import json
import socket
import shutil
import pycdlib
import gstage4
import wstage4
import pathlib
import tempfile
import subprocess
import grub_install
import strict_openssh
from fm_util import Util
from fm_util import CloudCacheGentoo
from fm_util import TmpMount
from fm_param import FmConst


class AssistOsBuilder:

    def __init__(self, infoPrinter, cloudOverlayDb, machine, bHasMemTest, showName):
        self._targetSystemInfo = {
            "amd64": {
                "subarch": "amd64",
                "variant": "systemd",
                "profile": "default/linux/amd64/23.0/no-multilib",
                "stage3-file": None,
            },
            "arm64": {
                "subarch": "arm64",
                "variant": "systemd",
                "profile": None,
                "stage3-file": None,
            },
        }
        self._p = infoPrinter
        self._cloudOverlayDb = cloudOverlayDb
        self._machine = machine
        self._bHasMemTest = bHasMemTest
        self._showName = showName
        self._snapshotFile = None

    def downloadFiles(self, bTitle):
        if bTitle:
            self._p.printInfo("- Downloading %s files..." % (self._showName))

        cache = CloudCacheGentoo(FmConst.gentooLinuxCacheDir)
        cache.sync()
        for arch, v in self._targetSystemInfo.items():
            assert arch in cache.get_arch_list()
            assert v["subarch"] in cache.get_subarch_list(arch)
        for arch, v in self._targetSystemInfo.items():
            try:
                v["stage3-file"] = cache.get_latest_stage3(arch, v["subarch"], v["variant"], cached_only=True)   # prefer local stage3 file
            except FileNotFoundError:
                v["stage3-file"] = cache.get_latest_stage3(arch, v["subarch"], v["variant"])
        self._snapshotFile = cache.get_latest_snapshot()                                                         # always use newest snapshot

    def build(self, arch, diskFsUuid, wdir):
        # for debugging purpose, use a local-op executable
        Util.cmdExec(os.path.join(FmConst.libexecDir, "lop-build-assist-os.py"),
                     str(self._p.getIndent()),
                     arch,
                     self._targetSystemInfo[arch]["profile"],
                     self._targetSystemInfo[arch]["stage3-file"][0],
                     self._targetSystemInfo[arch]["stage3-file"][1],
                     self._snapshotFile,
                     diskFsUuid,
                     wdir.path,
                     "1" if self._bHasMemTest else "0")

    def updateInitramfs(self, diskFsUuid, wdir):
        buf = """
            import bbki

            class MyConfig(bbki.ConfigBase):

                @property
                def flags(self):
                    return bbki.Flags.KERNEL | bbki.Flags.INITRAMFS

                def get_fs_layout(self):
                    return bbki.FsLayout()

                def get_build_variable(self, var_name):
                    return ""

                def get_mount_points(self):
                    return [
                        bbki.HostMountPoint(**{
                            "mountpoint": "/",
                            "dev_path_or_uuid": "UUID:FILE=%s:os/%s/rootfs.sqfs",
                            "fstype": "squashfs+overlay",
                            "opts": "",
                            "underlay_disk": {
                                "type": "image-file-on-removable",
                                "uuid": "UUID=%s",
                                "used_partition_table_types": ["mbr", "gpt"],
                                "image_file_path": "os/%s/rootfs.sqfs",
                                "children": [],
                            },
                        }),
                    ]

                def get_system_init(self):
                    return bbki.SystemInit(bbki.SystemInit.Type.SYSTEMD)

                def get_online_mount_points(self):
                    assert False

                def get_boot_mode(self):
                    assert False

                def get_bootloader_extra_time(self):
                    assert False

                def get_kernel_extra_init_cmdline(self):
                    assert False

            bbkiObj = bbki.Bbki(MyConfig())
            updater = bbkiObj.obtain_kernel_updater(bbkiObj.get_newest_boot_entry())
            updater.update_initramfs()
        """
        buf = re.sub("^            ", "", buf, flags=re.M)
        buf = buf % (diskFsUuid, wdir.get_arch(), diskFsUuid, wdir.get_arch())

        fullfn = os.path.join(wdir.get_latest_action_dirpath(), "root", "update-initramfs")
        pathlib.Path(fullfn).write_text(buf)
        with gstage4.Runner(wdir.get_arch(), wdir.get_latest_action_dirpath()) as c:
            c.shell_call("", "python /root/update-initramfs")
        os.unlink(fullfn)


class AssistDiskBuilder:

    def __init__(self, diskFsUuid, diskLabel, diskName):
        self._diskFsUuid = diskFsUuid
        self._diskLabel = diskLabel
        self._diskName = diskName
        self._wdirStage4GentooList = []
        self._wdirStage4WinXp = None
        self._wdirStage4Win7 = None
        self._targetSysDict = {}            # dict<arch, wdir>

    def addStage4(self, osType, wdir):
        if osType == "gentoo":
            assert isinstance(wdir, gstage4.WorkDir)
            assert wdir.is_build_finished()
            assert wdir.get_arch() not in [x.get_arch() for x in self._wdirStage4GentooList]
            self._wdirStage4GentooList.append(wdir)
        elif osType == "mswin_xp":
            assert isinstance(wdir, wstage4.WorkDir)
            assert wdir.is_build_finished()
            assert self._wdirStage4WinXp is None
            self._wdirStage4WinXp = wdir
        elif osType == "mswin_7":
            assert isinstance(wdir, wstage4.WorkDir)
            assert wdir.is_build_finished()
            assert self._wdirStage4Win7 is None
            self._wdirStage4Win7 = wdir
        else:
            assert False

    def addAssistOs(self, wdir):
        assert isinstance(wdir, gstage4.WorkDir) and wdir.is_build_finished()
        assert wdir.get_arch() not in self._targetSysDict
        self._targetSysDict[wdir.get_arch()] = wdir

    def exportToIsoFile(self, filePath):
        assert len(self._targetSysDict) > 0

        iso = pycdlib.PyCdlib()
        iso.new(udf="2.60")
        try:
            # add README.TXT
            buf = ""
            buf += 'This disc contains a "UDF" file system and requires an operating system\n'
            buf += 'that supports the ISO-13346 "UDF" file system specification.\n'
            buf = buf.encode("iso8859-1")
            iso.add_fp(io.ByteIO(buf), len(buf), iso_path="/README.TXT")

            # add files
            f = iso.get_udf_facade()
            f.add_directory("/os")
            f.add_directory("/data")

            # add boot files
            pass

            # write
            iso.write(filePath)
        finally:
            iso.close()

    def exportToDiskImageFile(self, filePath, imageSize):
        tDir = tempfile.mkdtemp()
        try:
            mntDir = os.path.join(tDir, "mnt")
            os.mkdir(mntDir)
            loopFile = os.path.join(tDir, "loop")

            Util.createAndMountFatDiskImage(filePath, imageSize, "mbr", loopFile, mntDir, fsUuid=self._diskFsUuid, label=self._diskLabel)
            try:
                mnt = grub_install.RootfsOrBootMountPoint(filePath, loopFile, mntDir, "vfat", "", False, True)
                self._exportToDir(mnt)
            finally:
                subprocess.check_call(["umount", mntDir])
                Util.loopUnMount(loopFile)
        finally:
            shutil.rmtree(tDir)

    def exportToUsbStick(self, devPath):
        assert "amd64" in self._targetSysDict or "arm64" in self._targetSysDict

        partDevPath = Util.formatDisk(devPath, "mbr", fsType="fat", fsUuid=self._diskFsUuid, label=self._diskLabel)
        with TmpMount(partDevPath) as mp:
            mnt = grub_install.RootfsOrBootMountPoint(devPath, partDevPath, mp.mountpoint, mp.fstype, mp.opts, False, True)
            self._exportToDir(mnt)

    def _exportToDir(self, mnt):
        # each volume must have size of 4G because we are using vfat file system
        # FIXME: 4G reports too large error, why?
        volumeSize = "2G"

        uuid = Util.getBlkDevUuid(mnt.partition)
        if uuid == "":
            raise Exception("can not get FS-UUID for %s" % (mnt.partition))

        osDir = os.path.join(mnt.mountpoint, "os")
        os.mkdir(osDir)
        dataDir = os.path.join(mnt.mountpoint, "data")
        os.mkdir(dataDir)

        # copy gentoo stage4 files
        for wdir in self._wdirStage4GentooList:
            dstFile = os.path.join(dataDir, "gentoo-%s.tar.xz" % (wdir.get_arch()))
            Util.compressMultiVolumeTarXz(wdir.get_latest_action_dirpath(), dstFile, volumeSize)

        # copy mswin_xp stage4 files
        if self._wdirStage4WinXp is not None:
            with self._wdirStage4WinXp.mount_image_file(stage=1) as m:
                Util.cloneNtfsMultiVolumeXz(m.drive_c_loop_file, os.path.join(dataDir, "mswin_xp-drive_c.img.xz"), [], volumeSize)

        # copy mswin_7 stage4 files
        if self._wdirStage4Win7 is not None:
            excludeList = [
                "Boot", "bootmgr",            # exclude boot files, we will recreate BIOS/EFI boot files when installing to target machine
                "BOOTSECT.BAK",               # exclude recovery file
            ]
            with self._wdirStage4Win7.mount_image_file(stage=1) as m:
                Util.cloneNtfsMultiVolumeXz(m.drive_c_loop_file, os.path.join(dataDir, "mswin_7-drive_c.img.xz"), excludeList, volumeSize)

        # install target system files (amd64) into usb stick
        if "amd64" in self._targetSysDict:
            sp = self._targetSysDict["amd64"].get_latest_action_dirpath()
            dstOsDir = os.path.join(osDir, "amd64")

            os.mkdir(dstOsDir)
            for fn in os.listdir(os.path.join(sp, "boot")):
                if fn.startswith("kernel-"):
                    shutil.copy(os.path.join(sp, "boot", fn), os.path.join(dstOsDir, "vmlinuz"))
                elif fn.startswith("initramfs-") and not fn.startswith("initramfs-files-"):
                    shutil.copy(os.path.join(sp, "boot", fn), os.path.join(dstOsDir, "initramfs.img"))
            shutil.copy(os.path.join(sp, "usr", "share", "memtest86+", "memtest.efi64"), dstOsDir)
            shutil.copy(os.path.join(sp, "usr", "share", "memtest86+", "memtest64.bios"), dstOsDir)
            Util.makeSquashedRootfsFiles(sp, dstOsDir, exclude=["boot/*"])

        # install target system files (arm64) into usb stick
        if "arm64" in self._targetSysDict:
            # FIXME
            assert False

        # install grub
        s = grub_install.Source(sp)
        t = grub_install.Target(grub_install.TargetType.MOUNTED_HDD_DEV, grub_install.TargetAccessMode.W, boot_mount_point=mnt)
        if "amd64" in self._targetSysDict:
            t.install_platform(grub_install.PlatformType.X86_64_EFI, s, removable=True, update_nvram=False)
            t.install_platform(grub_install.PlatformType.I386_PC, s, rs_codes=False)
        if "arm64" in self._targetSysDict:
            # FIXME
            assert False
        t.install_data_files(s, locales="*", fonts="*", themes="*")

        # create grub.cfg
        with open(os.path.join(mnt.mountpoint, "grub", "grub.cfg"), "w") as f:
            f.write("set default=1\n")
            f.write("set timeout=30\n")
            f.write("set gfxmode=auto\n")
            f.write("\n")
            f.write("insmod all_video\n")
            f.write("insmod videotest\n")
            f.write("insmod videoinfo\n")
            f.write("insmod gfxterm\n")
            f.write("terminal_output gfxterm\n")
            f.write("\n")

            if "amd64" in self._targetSysDict and "arm64" in self._targetSysDict:
                # FIXME
                assert False
            elif "amd64" in self._targetSysDict:
                f.write("menuentry \"Boot %s\" --class gnu-linux --class os {\n" % (self._diskName))
                f.write("    linux /os/amd64/vmlinuz\n")
                f.write("    initrd /os/amd64/initramfs.img\n")
                f.write("}\n")
                f.write("menuentry \"Boot %s (with Remote Control Agent)\" --class gnu-linux --class os {\n" % (self._diskName))
                f.write("    linux /os/amd64/vmlinuz rcagent\n")
                f.write("    initrd /os/amd64/initramfs.img\n")
                f.write("}\n")
            elif "arm64" in self._targetSysDict:
                # FIXME
                assert False
            else:
                assert False
            f.write("\n")

            if "amd64" in self._targetSysDict and "arm64" in self._targetSysDict:
                # FIXME
                assert False
            elif "amd64" in self._targetSysDict:
                f.write('if [ "${grub_platform}" == "efi" ]; then\n')
                f.write("    menuentry \"Run Memtest86+ (RAM test)\" {\n")
                f.write('        chainloader "/os/amd64/memtest.efi64"\n')
                f.write("    }\n")
                f.write('elif [ "${grub_platform}" == "pc"]; then\n')
                f.write("    menuentry \"Run Memtest86+ (RAM test)\" {\n")
                f.write('        linux "/os/amd64/memtest64.bios"\n')
                f.write("    }\n")
                f.write('fi\n')
            elif "arm64" in self._targetSysDict:
                # FIXME
                assert False
            else:
                assert False
            f.write("\n")

            # menuentry "Hardware Information (HDT)" {
            #     linux /os/%ARCH%/hdt
            # }

            # Menu
            f.write("menuentry \"Restart\" {\n")
            f.write("    reboot\n")
            f.write("}\n")
            f.write("\n")

            # Menu
            f.write("menuentry \"Power Off\" {\n")
            f.write("    halt\n")
            f.write("}\n")

        # create livecd
        # FIXME: it sucks that genkernel's initrd requires this file
        with open(os.path.join(mnt.mountpoint, "livecd"), "w") as f:
            f.write("")


class TargetFeatureAssistOs:

    def __init__(self, machine):
        self._machine = machine

    def update_target_settings(self, target_settings):
        assert "80-assist-os" not in target_settings.pkg_use_files

        target_settings.pkg_use_files["80-assist-os"] = self._useFileContent.strip("\n") + "\n"

    def get_custom_action_robustify_by_barkspider(self):
        custom_script_list = []

        # re-install packages so that they can be robustified
        custom_script_list.append(gstage4.scripts.OneLinerScript("emerge -1 dev-vcs/git net-misc/rsync net-misc/wget"))

        # change configure so that the applications can be robustified
        custom_script_list.append(gstage4.scripts.OneLinerScript("barkspider robustify all"))

        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_build_kernel_and_initramfs(self, arch, diskFsUuid, indentStr):
        jobcountMake, loadavg = Util.portageDetectJobCountAndLoadAverage(self._machine)

        buf = """
            import bbki
            import platform

            class MyConfig(bbki.ConfigBase):

                @property
                def flags(self):
                    return bbki.Flags.KERNEL | bbki.Flags.INITRAMFS

                def get_fs_layout(self):
                    return bbki.FsLayout()

                def get_build_variable(self, var_name):
                    return "--jobs=%d --load-average=%d"

                def get_mount_points(self):
                    return [
                        bbki.HostMountPoint(**{
                            "mountpoint": "/",
                            "dev_path_or_uuid": "UUID:FILE=%s:os/%s/rootfs.sqfs",
                            "fstype": "squashfs+overlay",
                            "opts": "",
                            "underlay_disk": {
                                "type": "image-file-on-removable",
                                "uuid": "UUID=%s",
                                "used_partition_table_types": ["mbr", "gpt"],
                                "image_file_path": "os/%s/rootfs.sqfs",
                                "children": [],
                            },
                        }),
                    ]

                def get_system_init(self):
                    return bbki.SystemInit(bbki.SystemInit.Type.SYSTEMD)

                def get_online_mount_points(self):
                    assert False

                def get_boot_mode(self):
                    assert False

                def get_bootloader_extra_time(self):
                    assert False

                def get_kernel_extra_init_cmdline(self):
                    assert False

            indentStr = "%s"

            bbkiObj = bbki.Bbki(MyConfig())

            bootEntrySpec = bbki.BootEntrySpec({
                "arch": platform.machine(),
                "kernel_type": "linux",
                "kernel_params": {
                },
                "kernel_addon_params": {
                    "pnp": {
                        "peripheral": None,
                    },
                    "noatime_default": {
                    },
                    "wireless_regdb": {
                    },
                    "linux_firmware": {
                        "strict": None
                    },
                    "fedora_alsa_firmware": {
                        "strict": None
                    },
                    "riptide_firmware": {
                        "strict": None
                    },
                    "ipw2100_firmware": {
                        "strict": None
                    },
                    "ipw2200_firmware": {
                        "strict": None
                    },
                    "storage_layouts": {
                        "efi-bcache-btrfs": None,       # for target os
                        "efi-btrfs": None,              # for target os
                        "efi-ext4": None,               # for target os
                        "efi-msr-ntfs": None,           # for target os
                        "bios-ext4": None,              # for target os
                        "bios-ntfs": None,              # for target os
                        "squashfs-on-fat": None,
                        "squashfs-on-isofs": None,
                    },
                    "x86_32": {
                    },
                    "x86_intel": {
                    },
                    "x86_amd": {
                    },
                    "iwd": {
                    },
                    "vm_as_host": {
                    },
                    "vt": {
                    },
                    # "display_full": {
                    # },
                    # "display_simple": {
                    #     "boot-logger": None,
                    # },
                    "desktop_functions": {
                    },
                },
                "initramfs_params": {
                },
            })

            print(indentStr, end="")
            print("        - Synchronizing...")
            bbkiObj.sync(bootEntrySpec)

            print(indentStr, end="")
            print("        - Fetching...")
            bbkiKernelBuilder = bbkiObj.obtain_kernel_installer(bootEntrySpec)
            bbkiKernelBuilder.fetch_kernel()
            for name in bbkiKernelBuilder.spec.kernel_addons:
                bbkiKernelBuilder.fetch_kernel_addon(name)

            print(indentStr, end="")
            print("        - Extracting...")
            bbkiKernelBuilder.unpack()

            print(indentStr, end="")
            print("        - Patching...")
            bbkiKernelBuilder.patch_kernel()

            print(indentStr, end="")
            print("        - Generating .config file...")
            bbkiKernelBuilder.generate_kernel_config_file()

            print(indentStr, end="")
            print("        - Building and installing...")
            bbkiKernelBuilder.install_kernel()

            print(indentStr, end="")
            print("        - Creating initramfs...")
            bbkiKernelBuilder.install_initramfs()
        """
        buf = re.sub("^            ", "", buf, flags=re.M)
        buf = buf % (jobcountMake, loadavg, diskFsUuid, arch, diskFsUuid, arch, indentStr)

        custom_script_list = []
        custom_script_list.append(gstage4.scripts.CreateFileFunction("/root/build-kernel", buf))
        custom_script_list.append(gstage4.scripts.OneLinerScript("python /root/build-kernel"))
        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_generate_machine_id(self):
        s = gstage4.scripts.CreateFileFunction("/etc/machine-id", Util.generateRandomMachineId() + "\n")
        return gstage4.custom_actions.SimpleCustomAction(s)

    # def get_custom_action_establish_seat_main(self):
    #     s = gstage4.scripts.PlacingFilesScript()
    #     s.append_file("/etc/udev/rules.d/72-seat-main.rules", 'TAG=="seat", ENV{ID_SEAT}="seat-main"\n')
    #     return gstage4.custom_actions.SimpleCustomAction(s)

    def get_custom_action_enable_console_service(self):
        buf = ""
        buf += "[Unit]\n"
        buf += "Description=Assistant Disk Console\n"
        buf += "\n"
        buf += "[Service]\n"
        buf += "ExecStart=/usr/bin/cage -- /usr/bin/kitty\n"
        buf += "\n"
        buf += "[Install]\n"
        buf += "WantedBy=graphical.target\n"

        s = gstage4.scripts.CreateFileFunction("/etc/systemd/system/assist-disk-console.service", buf)
        return gstage4.custom_actions.SimpleCustomAction(s)

    def get_custom_action_config_services(self):
        custom_script_list = []

        # sshd
        class ModifySshdConfig(gstage4.FunctionInChroot):
            def execute(self, chroot_dir_hostpath):
                obj = strict_openssh.ServerSysCfg(prefix=chroot_dir_hostpath)
                obj.generate_host_key("rsa")
                obj.generate_host_key("ecdsa")
                # obj.generate_host_key("ed25519")      # FIXME paramiko does not support generating this key
                obj.set_using_pam(True)
                obj.set_permit_root_login("yes")
                obj.accept_common_environment_variables()
        custom_script_list.append(ModifySshdConfig())

        # NetworkManager
        class ModifyNetworkManagerConfig(gstage4.FunctionInChroot):
            def execute(self, chroot_dir_hostpath):
                fullfn = os.path.join(chroot_dir_hostpath, "usr", "lib", "systemd", "system", "NetworkManager.service")
                buf = pathlib.Path(fullfn).read_text()
                buf = re.sub(r"^Description=.*$", r"\g<0>\nConditionKernelCommandLine=!rcagent", buf, flags=re.M)
                pathlib.Path(fullfn).write_text(buf)
        custom_script_list.append(ModifyNetworkManagerConfig())

        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_config_rcagent(self):
        custom_script_list = []

        argData = {
            "bluetooth_allowed_mac_addresses": [],                          # FIXME
            "bluetooth_pin": Util.generateRandomDigits(4),                  # FIXME
            "bluetooth_passkey": Util.generateRandomDigits(6),              # FIXME
            "wifi_ssid": FmConst.rcWifiSsid,
            "wifi_password": Util.generateRandomDigits(8),                  # FIXME
        }
        custom_script_list.append(gstage4.scripts.OneLinerScript("/usr/libexec/rcagent-generate-conf '%s'" % (json.dumps(argData))))

        return gstage4.custom_actions.SimpleCustomAction(*custom_script_list)

    def get_custom_action_create_livecd_mount_point(self):
        assert False

    _useFileContent = """
"""


class AssistDiskRemote:

    def __init__(self, hostname, port):
        self.hostname = hostname
        self.ctrlPort = port

        self.sock = None
        self.sslSock = None

    def connectAndInit(self):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.connect((self.hostname, self.ctrlPort))

            ctx = SSL.Context(SSL.TLSv1_2_METHOD)
            ctx.use_certificate_file(FmConst.myCertFile)
            ctx.use_privatekey_file(FmConst.myPrivKeyFile)
            self.sslSock = SSL.Connection(ctx, self.sock)
            self.sslSock.set_connect_state()

            self._sendRequestObj({
                "command": "init",
                "hostname": socket.gethostname(),
                "cpu-arch": Util.getCpuArch(),
                "cpu-model": Util.getCpuModel(),
                "plugin": "gentoo",
            })
            resp = self._recvReponseObj()
            if "error" in resp:
                raise Exception(resp["error"])
        except:
            self.dispose()
            raise

    def dispose(self):
        assert self.asyncJobCount == 0

        if self.identityFile is None:
            os.unlink(self.identityFile)
            self.identityFile = None
        if self.cfgFile is None:
            os.unlink(self.cfgFile)
            self.cfgFile = None
        self.wSshPort = None
        self.wRsyncPort = None
        self.wCatFilePort = None

        if self.sslSock is not None:
            self.sslSock.close()
            self.sslSock = None
        if self.sock is not None:
            self.sock.close()
            self.sock = None

        self.ctrlPort = None
        self.hostname = None

    def createStorageLayout(self, layoutName):
        pass

    def sendSystemFiles(self, layoutName):
        pass
